# test_sql_generator.py
import sys
import os

# 添加项目路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from agents.sql_generator import SQLGenerationAgent
from models.state import WorkflowState

def test_sql_generation():
    """测试SQL生成Agent"""
    agent = SQLGenerationAgent()
    
    # 测试用例 - 模拟不同意图的输入
    test_cases = [
        {
            "input": "显示去年的销售总额",
            "intent": "Imperativequery",
            "entities": {"tables": ["sales"], "columns": ["amount"], "time": "last_year"}
        },
        {
            "input": "找出最畅销的产品", 
            "intent": "ExploratoryAnalysis",
            "entities": {"tables": ["sales", "products"], "columns": ["product_name", "amount"]}
        },
        {
            "input": "导出所有客户信息",
            "intent": "Dataextraction", 
            "entities": {"tables": ["users"], "columns": ["id", "username", "email"]}
        },
        {
            "input": "分析各区域的销售趋势",
            "intent": "ExploratoryAnalysis",
            "entities": {"tables": ["sales"], "columns": ["region", "amount", "sale_date"]}
        }
    ]
    
    print("=== SQL生成测试 ===")
    
    for i, test_case in enumerate(test_cases, 1):
        print(f"\n{'='*50}")
        print(f"测试用例 {i}: {test_case['input']}")
        print(f"意图类型: {test_case['intent']}")
        
        # 创建状态对象
        state = WorkflowState(
            user_input=test_case['input'],
            intent=test_case['intent'],
            entities=test_case['entities'],
            schema_info={"tables": ["sales", "users", "products"]}
        )
        
        # 生成SQL
        result_state = agent.process(state)
        
        print(f"生成的SQL: {result_state.generated_sql}")

def interactive_sql_generation():
    """交互式SQL生成测试"""
    agent = SQLGenerationAgent()
    
    print("=== 交互式SQL生成测试 ===")
    print("输入问题，系统将生成相应的SQL语句")
    print("输入 '退出' 结束程序")
    print("-" * 50)
    
    while True:
        user_input = input("\n请输入问题: ").strip()
        
        if user_input in ['退出', 'quit', 'exit']:
            print("再见！")
            break
            
        if not user_input:
            continue
            
        # 先进行意图识别（简化版）
        intent = "Imperativequery"  # 默认，实际应该调用意图识别Agent
        if any(word in user_input for word in ['分析', '趋势', '找出', '比较']):
            intent = "ExploratoryAnalysis"
        elif any(word in user_input for word in ['导出', '列出', '所有', '详细']):
            intent = "Dataextraction"
        
        # 创建状态对象
        state = WorkflowState(
            user_input=user_input,
            intent=intent,
            entities={},
            schema_info={"tables": ["sales", "users", "products"]}
        )
        
        # 生成SQL
        result_state = agent.process(state)
        
        print(f"识别意图: {intent}")
        print(f"生成的SQL:")
        print(f"   {result_state.generated_sql}")

if __name__ == "__main__":
    # 选择测试模式
    print("请选择测试模式:")
    print("1. 批量测试")
    print("2. 交互式测试")
    
    choice = input("请输入选择 (1 或 2): ").strip()
    
    if choice == "1":
        test_sql_generation()
    else:
        interactive_sql_generation()